feat(scheduler): range-repartition rule primitives (batteries-included) - #2196
Open
avantgardnerio wants to merge 25 commits into
Open
feat(scheduler): range-repartition rule primitives (batteries-included)#2196avantgardnerio wants to merge 25 commits into
avantgardnerio wants to merge 25 commits into
Conversation
Scheduler-side AQE machinery so any optimizer rule that inserts an `UnorderedRangeRepartitionExec` or `OrderedRangeRepartitionExec` gets end-to-end correct routing for free. A range-partitioning rule now needs to do only one thing: splice a `RuntimeStatsExec` + URRE/ORRE above whatever it wants range-repartitioned. Everything downstream is handled: - Executor already ships per-sub-part quantile sketches to the scheduler (apache#2094 / apache#2175). - Scheduler recognizes the range repartition at stage-completion time, merges sketches into `K − 1` global cuts, and rewrites the downstream location map so partition `k` pulls from every producer file whose sketched `[min, max]` overlaps `k`'s assigned cut range. - Cuts + routing expression are parked on the boundary `ExchangeExec`; at task-specialization time the adapter wraps the `ShuffleReader` in a `PerPartitionFilterExec` (apache#2195) that trims straddling sub-parts to their assigned slice. Primitives a rule writer now has: - `plan_contains_range_repartition` / `find_range_repartition_routing_expr` (runtime_stats.rs) — detect URRE or ORRE anywhere in a plan subtree. - `compute_overlapping_locations` / `overlap_remap_partitions` (runtime_stats.rs) — merged sketches + cuts → per-downstream-partition location map with correct producer-file assignment. - `ExchangeExec.range_repartition_routing` slot + `resolve_*` / `.range_repartition_routing()` accessors — scheduler → adapter handoff. Mirrors the existing `coalesce: Arc<Mutex<Option<Arc<CoalescePlan>>>>` slot pattern (same lifecycle, same `with_new_children` carry-through, same idempotent-overwrite semantics — see `b839f036` /`232d7611` for the coalesce refactors that established the pattern). - DER classifier arm (distributed_exchange.rs) — recognizes `RuntimeStatsExec → URRE/ORRE` as a stage boundary so the sketches ship to the scheduler at stage-N completion. All of this fires only when a rule has actually inserted an RRE. On plans without one, `plan_contains_range_repartition` returns `false` and the machinery no-ops — zero overhead for the 99%+ non-RRE case. Test coverage (15 new unit tests, all pure-function): - `plan_walker_tests` (6): URRE and ORRE at root and nested, bare source returns `None`/`false`. - `overlap_remap_tests` (5): disjoint, straddling, missing-file_id error, default/report mismatch error, empty-sketch passthrough. - `range_repartition_routing_tests` (4): unresolved returns `None`, resolve→get roundtrip, second resolve overwrites, `with_new_children` preserves the slot. Full end-to-end wiring is exercised in the follow-up consumer rule PRs (AdaptiveRangeShuffleRule for join-side rewrites, ParallelWindowDetectRule for parallel windows) where a live URRE-inserting fixture is naturally available; deferring the ~300 LOC of AQE-graph scaffolding here in favor of pure-function coverage. Existing tests: 222 in ballista-core, 281 in ballista-scheduler — all pass. Clippy clean, fmt clean. Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
Contributor
Author
|
@phillipleblanc this one is the base upon which anyone can build a range-repartitioned optimizer rule. I intend to use it for parallel windows in the next PR, but I was also playing with it to prevent spilling in SortShuffle on TPC-H q20. It's interesting because now we can play around and try things out. |
Rustdoc CI (`RUSTDOCFLAGS=-D warnings`) rejects public docs linking to private items. `range_repartition_common` is `mod` (not `pub mod`) and `split_batch_by_range` is `pub(super)`, so neither is reachable from the public rustdoc surface for `range_partition_predicates`. Convert the intra-doc link to plain text; the code cross-reference still reads for a curious reviewer grepping the tree. Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
Contributor
|
I won't have time to review this in depth for a couple of days, but I will look at it sometime this week. |
…mpotency The DER classifier's fourth arm inserts a passthrough `ExchangeExec(None)` above a `RuntimeStatsExec → URRE/ORRE` chain top, and its docstring claims a second pass over the resulting plan is a no-op — the freshly-inserted Exchange is caught by the `.downcast_ref::<ExchangeExec>().is_none()` outer guard, and from any upstream node's perspective an `ExchangeExec` is not a chain top (only `RuntimeStatsExec` is). Both halves of that claim are load-bearing (a bug in either would double-wrap the chain top on every AQE replan) but were only covered by the follow-up consumer-rule PRs. Two direct tests here: - `range_repartition_chain_top_gets_exchange_inserted` — verifies that arm 4 fires on `RoundRobin → RSE → URRE → leaf` and inserts the Exchange between the parent and the chain top, leaving the chain itself intact underneath. - `range_repartition_chain_top_insertion_is_idempotent` — runs the rule twice against the same starting plan and asserts the second pass adds no additional Exchange and produces an identical plan text. Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
`i` and `k` say nothing about what they index. In these helpers both are partition-space indices, so name them for what they are: - `k` → `partition_count` (in `range_partition_predicates`, `compute_overlapping_locations`) - `|i|` closure param → `|partition_idx|`; the inner `|j|` cut lookup → `|cut_idx|` - `partition_k` → `partition_idx` (matches the closure param above) - `out` → `overlaps_by_partition` (matches what the return doc calls it) - `smin` / `smax` → `sketch_min` / `sketch_max` - `spl` → `assignment` (it's one element of `assignments`) - `inner` → `remapped_bucket` - Plan walkers' `|c|` / `for c in ...` → `|child|` / `for child in ...` `lit`, `ge`, `lt`, `lo`, `hi` left as-is — those are established domain abbreviations, not opaque one-letter loop vars. Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
- Cut the 17-line "Fourth classifier arm" block down to the one thing the code doesn't already say: why we insert a boundary here at all (executor's report walker collects sketches at stage boundaries). - Drop "arm 4 / arm 3" numbering in the test-section header, test docstrings, and assertion messages — positional numbering rots the moment someone reorders the classifier's else-if chain or adds a case. - Rename the fourth arm's `|c|` closure params to `|child|`. - Fully railroad `is_range_repartition_chain_top` — every branch has an explicit early return, no `||` chain at the tail. Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
…n metadata Comment: - Cut the 13-line block above the PerPartitionFilterExec injection down to the one non-obvious why: straddling sub-parts + FinalPartitioned splitting partial sums. Log levels: - Downgrade the "injecting PerPartitionFilterExec" adapter trace, the "no non-empty sketches" fallback, and the happy-path overlap-remap summary from info! to debug!. All are per-range-repartition-stage events, useful for debug but not for baseline INFO output. Two of Brent's TODOs on the AQE hook flagged silent-degradation paths: - Empty `runtime_stats_reports` for a stage the planner detected as a range repartition: passthrough routing would silently misroute straddling rows. Now returns `Err` — no sketches means no cuts and we can't route. - `plan_contains_range_repartition` returns true but the routing-expr walker returns None: previously logged and continued with `routing = None`, which skipped the downstream filter injection and produced wrong aggregates. Now returns `Err` — this is a plan-shape invariant violation, not a soft-fallback case. Both errors propagate through `update_stage_progress` and fail the job rather than produce incorrect results. Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
Bring RuntimeStatsReport, RuntimeStatsPartitionEntry, and QuantileSketchState into runtime_stats.rs at file scope so signatures for merge_reports and kin drop the crate::serde::protobuf:: prefix. Bring the range-repartition helpers into aqe/mod.rs the same way; rename single-letter closure vars (t/m/c/e/v) and trim the error and debug messages to one line each. Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
finder, rename overlap-remap arg to original_partitions Split plan_contains_range_repartition into two early-return ifs plus the trailing recursive .any(). Switch find_range_repartition_routing_expr from Option to Result: the outer public fn returns Result<Arc<dyn PhysicalExpr>>; recursion moves into a private find_routing_expr returning Result<Option<_>>, so "not found in this subtree" (Ok(None)) and "found but order_by is empty" (Err) stay distinguishable. Match-on-slice replaces the silent order_by()[0] panics. Rename overlap_remap_partitions' `default` argument (and the scheduler's local) to `original_partitions` — the semantics are "before remap", not "fallback default". Import PhysicalExpr and PartitionLocation into runtime_stats.rs to keep the affected signatures on one line. Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
… to cut_partitions Break maybe_range_repartition_overlap_remap into two orthogonal pieces: range_repartition_routing (recover cuts + routing expression from merged sketches; Ok(None) for legit passthrough cases, Err for invariant breaks) and the caller composing that with cut_partitions (previously overlap_remap_partitions). The old fn conflated plan-shape detection, sketch merging, and partition remap; now the caller uses plan_contains_range_repartition as a guard and slots in the routing where it belongs. Distinguish two "empty cuts" cases that used to collapse into one silent debug! + passthrough: total_rows == 0 (no data, passthrough is safe) and partition_count < 2 (K=1, single-partition range-repartition is degenerate but valid) both return Ok(None); any other empty-cuts case is a real sketch bug and now errors out. Also match on merged.as_slice() rather than merged.first() to catch the 2+ groups shape bug that first() silently ignored. Rename overlap_remap_partitions -> cut_partitions to match the "cut the K-space by these boundaries" mental model. Terse one-liner error messages throughout. Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
plan_contains_range_repartition and find_range_repartition_routing_expr walked the same tree looking for URRE/ORRE, called back-to-back at the same call site. Merge into range_repartition_routing_expr returning Result<Option<Arc<dyn PhysicalExpr>>>: Ok(None) means no operator in the plan (single-mode absence), Err means one was found but its order_by was empty (invariant break). Caller uses .is_some() for the guard in the let-chain and passes the recovered expr into range_repartition_routing so the routing recovery no longer walks the plan itself. Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
Sweep the sites this PR introduced (across shuffle_writer.rs and distributed_exchange.rs) that reached for the boolean via the longer form. DataFusion provides is::<T: ExecutionPlan>() directly on dyn ExecutionPlan alongside downcast_ref, so the terse form is the right fit when we only need the check. Where the negation reads better as !.is::<T>() it replaces .downcast_ref().is_none(). Collapses several 6-line if-blocks (imports for UnorderedRangeRepartitionExec / OrderedRangeRepartitionExec at the top of shuffle_writer.rs let the ||-chain fit on one line) and lets the symmetric arms of is_range_repartition_chain_top stay as two one-liners each. Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
…ent skip Drive cut_partitions from original_partitions directly; look up sketches by (task_id, sub_part_id) from a small map built once out of the reports. Deletes compute_overlapping_locations and SubPartLocation — the intermediate identity vector isn't needed once the walk is driven from the physical records. Files without a usable sketch (missing entry, or count == 0) are safe to skip only when partition_stats.num_rows == Some(0). Some(n > 0) or None now surface as an error rather than silently dropping data. Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
…s per file The inner K-loop over `remapped` made the walk O(F·K) — quadratic when the file count and output partition count are of similar scale (typical for a range shuffle). Since `global_cuts` is monotone and each file's sketch is a single [min, max] interval, the set of overlapping buckets is contiguous. Two `partition_point` calls locate `[b_lo, b_hi]` and we push into that slice, taking each file to O(log K + overlap). Adds a debug_assert on cuts monotonicity and a multi-cut test that exercises the range-walk boundaries (existing tests only had a single cut and couldn't catch an off-by-one). Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
Member
|
Really like the shape of this, and the "a rule only has to splice two operators" goal is the right one. I dug into one thing and it looks like a real gap, so I built a repro rather than just asserting it. The issue
let remapped = cut_partitions(partitions, reports, &routing.cuts)?; // straddlers duplicated
self.planner.set_repartition_routing(stage_id, routing)?; // may silently no-op
self.output_locations = partitions.into_iter().flatten().collect();The reachable shape is the canonical oneI initially assumed you would need something exotic to hit this. You do not. It is the exact
ReproThree tests against Leg 1, dropped into the #[test]
fn repro_root_level_chain_top_gets_no_exchange() {
let rule = DistributedExchangeRule::default();
let root = stats_over_urre_over_leaf(); // your existing helper
let result = rule.optimize(root, &config()).unwrap();
let adaptive = result.downcast_ref::<AdaptiveDatafusionExec>().unwrap();
assert!(adaptive.input().is::<RuntimeStatsExec>());
assert_eq!(count_exchanges(result.as_ref()), 0);
}Leg 2, composing No Leg 3, positive control so this is not just me misreading the mechanism. Same splice, same calls, only a Exchange inserted, cuts parked, everything works. Position in the plan is the only variable between leg 2 and leg 3. Full leg 2 and leg 3 test fileAdd use crate::state::aqe::execution_plan::RangeRepartitionRouting;
use crate::state::aqe::planner::AdaptivePlanner;
use ballista_core::execution_plans::{
RuntimeStatsExec, UnorderedRangeRepartitionExec, cut_partitions,
};
use ballista_core::extension::SessionConfigExt;
use ballista_core::serde::protobuf::{RuntimeStatsPartitionEntry, RuntimeStatsReport};
use ballista_core::serde::scheduler::{
ExecutorMetadata, ExecutorOperatingSystemSpecification, ExecutorSpecification,
PartitionId, PartitionLocation, PartitionStats,
};
use datafusion::arrow::compute::SortOptions;
use datafusion::arrow::datatypes::{DataType, Field, Schema};
use datafusion::datasource::memory::MemorySourceConfig;
use datafusion::datasource::source::DataSourceExec;
use datafusion::physical_expr::PhysicalSortExpr;
use datafusion::physical_plan::ExecutionPlan;
use datafusion::physical_plan::expressions::col;
use datafusion::prelude::SessionConfig;
use std::sync::Arc;
fn v_schema() -> Arc<Schema> {
Arc::new(Schema::new(vec![Field::new("v", DataType::Float64, false)]))
}
/// The canonical splice a range-repartition rule is told to emit,
/// here at the root of the plan.
fn stats_over_urre_root() -> Arc<dyn ExecutionPlan> {
let schema = v_schema();
let source: Arc<dyn ExecutionPlan> = Arc::new(DataSourceExec::new(Arc::new(
MemorySourceConfig::try_new(&[vec![]], schema.clone(), None).unwrap(),
)));
let sort_expr = PhysicalSortExpr {
expr: col("v", schema.as_ref()).unwrap(),
options: SortOptions { descending: false, nulls_first: false },
};
let urre: Arc<dyn ExecutionPlan> = Arc::new(
UnorderedRangeRepartitionExec::try_new(source, vec![sort_expr.clone()], 2).unwrap(),
);
Arc::new(RuntimeStatsExec::try_new(urre, Some(vec![sort_expr])).unwrap())
}
/// Same splice with an ordinary parent above it. A FilterExec survives the
/// optimizer pipeline. RoundRobin RepartitionExec and CoalescePartitionsExec
/// do NOT, both get stripped.
fn stats_over_urre_with_parent() -> Arc<dyn ExecutionPlan> {
let pred: Arc<dyn datafusion::physical_expr::PhysicalExpr> =
Arc::new(datafusion::physical_expr::expressions::BinaryExpr::new(
col("v", v_schema().as_ref()).unwrap(),
datafusion::logical_expr::Operator::Gt,
Arc::new(datafusion::physical_expr::expressions::Literal::new(
datafusion::scalar::ScalarValue::Float64(Some(0.0)),
)),
));
Arc::new(
datafusion::physical_plan::filter::FilterExec::try_new(pred, stats_over_urre_root())
.unwrap(),
)
}
fn location(sub_part_id: usize, producer_task_id: usize, rows: u64) -> PartitionLocation {
PartitionLocation {
map_partition_id: 0,
partition_id: PartitionId {
job_id: "repro-job".into(),
stage_id: 0,
partition_id: sub_part_id,
},
executor_meta: ExecutorMetadata {
id: format!("exec-{producer_task_id}"),
host: "".to_string(),
port: 0,
grpc_port: 0,
specification: ExecutorSpecification::default().with_vcores(0),
os_info: ExecutorOperatingSystemSpecification::default(),
},
partition_stats: PartitionStats::new(Some(rows), None, None),
file_id: Some(producer_task_id as u64),
is_sort_shuffle: false,
}
}
#[tokio::test]
async fn repro_routing_dropped_after_partitions_already_duplicated()
-> datafusion::error::Result<()> {
let config = SessionConfig::new_with_ballista();
let mut planner =
AdaptivePlanner::try_from_plan(&config, stats_over_urre_root(), "repro-job".into())?;
let stages = planner.runnable_stages()?;
let stage_id = stages
.as_ref()
.and_then(|s| s.first())
.map(|e| e.plan.stage_id())
.expect("a runnable stage must exist");
// Step 1: the remap, exactly as update_stage_progress calls it.
// Producer task 7, sub-part 0, sketched [5, 25], straddles the cut at 15.
let reports = vec![ballista_core::execution_plans::TaskRuntimeStats {
producer_task_id: 7,
report: RuntimeStatsReport {
order_by: vec![],
partitions: vec![RuntimeStatsPartitionEntry {
partition_id: 0,
row_count: 3,
sketch: Some(ballista_core::execution_plans::sketch_to_proto(
&datafusion_functions_aggregate_common::tdigest::TDigest::new(100)
.merge_unsorted_f64(vec![5.0, 15.0, 25.0]),
)?),
}],
},
}];
let cuts = vec![15.0];
let remapped = cut_partitions(vec![vec![location(0, 7, 3)]], &reports, &cuts)?;
assert_eq!(remapped[0].len(), 1, "straddler duplicated into partition 0");
assert_eq!(remapped[1].len(), 1, "straddler duplicated into partition 1");
// Step 2: park the routing so a read-side filter can trim it.
let routing = RangeRepartitionRouting {
cuts: cuts.clone(),
routing_expr: col("v", v_schema().as_ref()).unwrap(),
};
let parked = planner.set_repartition_routing(stage_id as usize, routing);
assert!(parked.is_ok(), "silently succeeds even though nothing was parked");
// Step 3: confirm nothing was actually parked.
let plan_str = format!(
"{}",
datafusion::physical_plan::displayable(planner.current_plan()).indent(true)
);
assert!(
!plan_str.contains("range_repartition_cuts"),
"no ExchangeExec carries the cuts, so no PerPartitionFilterExec downstream"
);
// What update_stage_progress assigns to output_locations.
let file_ids: Vec<u64> = remapped
.into_iter()
.flatten()
.map(|l| l.file_id.unwrap())
.collect();
assert_eq!(file_ids, vec![7, 7], "same file twice, 3 rows read as 6");
Ok(())
}
#[tokio::test]
async fn control_routing_is_parked_when_chain_top_has_a_parent()
-> datafusion::error::Result<()> {
let config = SessionConfig::new_with_ballista();
let mut planner = AdaptivePlanner::try_from_plan(
&config,
stats_over_urre_with_parent(),
"control-job".into(),
)?;
let stages = planner.runnable_stages()?;
let stage_id = stages
.as_ref()
.and_then(|s| s.first())
.map(|e| e.plan.stage_id())
.expect("a runnable stage must exist");
let routing = RangeRepartitionRouting {
cuts: vec![15.0],
routing_expr: col("v", v_schema().as_ref()).unwrap(),
};
planner.set_repartition_routing(stage_id as usize, routing)?;
let plan_str = format!(
"{}",
datafusion::physical_plan::displayable(planner.current_plan()).indent(true)
);
assert!(
plan_str.contains("range_repartition_cuts=1"),
"control: the cuts must be parked on the boundary ExchangeExec"
);
Ok(())
}A side finding that might matter moreI first wrote the control with a That is worth flagging on its own. What I did not proveI composed Suggested fixMake Nothing else in the review turned up anything. No public API breaks, the Investigation and repro here were LLM assisted, with the test output and plan dumps above coming from actual runs against the PR head. |
Two failing tests documenting the routing-park gap Andy flagged in PR apache#2196 review: when an RSE → URRE splice sits at the plan root, DER never wraps an ExchangeExec above it (transform_up only inspects children) and set_repartition_routing has nothing to park cuts on, so cut_partitions' straddler duplication reaches output_locations unfiltered. Both new tests assert the corrected behavior — exchange inserted, cuts parked — so they go red on current main and will go green when the fix lands. The FilterExec-parent sibling stays green as the positive control. Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
…park on The routing park silently returned Ok(()) when the stage's cached boundary wasn't an ExchangeExec, leaving cut_partitions' straddler duplication in place with no downstream filter to trim it — the silent outlier in a code path that otherwise fails loud on missing file_id, absent sketch, or bad merge_reports counts. Fail loud with a diagnostic error naming the actual boundary type instead. The one reachable trigger for this today (a range-repartition splice at the plan root) still leaves the URRE-at-root regression tests red; the DER-side fix that closes that shape lands next. Adds a synthetic bare-leaf regression test that pins the contract independent of DER: whenever the stage cache holds something other than an ExchangeExec, calling set_repartition_routing must error. Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
…e whitelist `DistributedExchangeRule::optimize` walks `transform_up`, which never visits the plan root as a child. A range-repartitioned splice at the root therefore got no boundary exchange wrapped above it, so `set_repartition_routing` had nothing to park cuts on. Wrap the root as a post-walk step, before the outer `AdaptiveDatafusionExec` goes on. Also splits the plan-shape whitelist out into a shared `plan_algebra` module: `preserves_distribution` stays where it was semantically (used by the sketch walker to descend past passthrough operators), and a looser `preserves_partitioning` sibling covers cases where only partition boundaries need to survive (rows and values within a partition are fair game). `is_range_repartitioned` drops its hardcoded `RuntimeStatsExec` check in favor of `preserves_partitioning` — the sketch tap is one of several partitioning-preserving nodes that can sit above a URRE/ORRE, and DF doesn't expose these as algebraic properties on `ExecutionPlan`. Flips the two red URRE-at-plan-root regression tests green; existing parent-based, idempotency, and control tests keep passing. Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
…ency at root Andy's PR apache#2196 review flagged that the two DER-only tests exercised a shape the real optimizer pipeline strips: a `RoundRobinBatch` `RepartitionExec` parent that survives `rule.optimize()` in isolation but gets removed by other passes before DER runs in the pipeline. So the tests passed on a plan shape production never produces, and a rule author could write something that looks protected and lose the protection at runtime. - Delete `range_repartition_chain_top_gets_exchange_inserted`. The parent case is already covered end-to-end by `routing_parks_when_range_repartition_has_a_parent`, which uses a `FilterExec` parent (survives the pipeline) and drives through `AdaptivePlanner` instead of calling the rule alone. - Rewrite the idempotency test to use the root-level shape directly (`stats_over_urre_over_leaf()` with no artificial parent) — this is a shape production actually produces, and running `rule.optimize()` twice on it still validates DER's short-circuit for `AdaptiveDatafusionExec` inputs. Also drops the last of the "chain top" jargon from test names and docstrings. Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
… optimize Two orthogonal cleanups on the DER rule: - Collapse the range-repartition arm to the same shape as Coalesce/SPM: identify pattern, grab the (single) child, wrap in `ExchangeExec`, return via `with_new_children`. The old `.any()` + `.map()` over children was defending against a multi-child parent case that can't occur — URRE only ever sits below a stage boundary, so its parent is unary in practice. - Flatten `optimize()`: early-return for the idempotent `AdaptiveDatafusionExec` input case so the root-wrap logic reads as a linear sequence rather than a nested `if/else`. Adds a TODO on the range-repart arm noting the intended endgame — kill it entirely when `ExchangeExec` carries a range-cuts partitioning variant the way it carries `Partitioning::Hash`, at which point URRE can replace itself the way the Hash arm above already does. Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
Contributor
Author
|
@andygrove thank you for your thorough review. I'm very happy with the review quality in this project. The feedback is always substantive and positive. I believe I addressed your points, but I did leave a TODO. Unfortunately we're waiting on Datafusion for that enum variant. (RangeRepartitioning) |
The previous commit put the TODO on the `transform()` range-repart arm, but the arm the user wanted killed is the root-wrap block in `optimize()` — that's the one that goes away when we wrap in `AdaptiveDatafusionExec` up front and let `transform_up` visit the plan root as a child. Moving the TODO there with the specific unblock path (mechanical plan_id snapshot updates). Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
Correct my earlier reasoning: the wrap-first-in-AdaptiveDatafusionExec approach isn't the right endgame — it shifts plan_id numbering because the wrapper gets allocated before the inner ops instead of after. The actual endgame is the same as the `transform()` range-repart arm: `ExchangeExec` learning a range-cuts partitioning variant like it already carries `Partitioning::Hash`. Once URRE can replace itself with an ExchangeExec the way the Hash arm does, `transform_up`'s output already has an ExchangeExec at the range-repart position, the plain `AdaptiveDatafusionExec` wrap handles the root case with no ceremony, plan_id numbering stays unchanged, and both this root-wrap branch and the range-repart arm in `transform()` above collapse together. Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
andygrove
reviewed
Aug 1, 2026
| return Ok(Transformed::yes(Arc::new(exchange_exec))); | ||
| } | ||
| } else if !execution_plan.is::<ExchangeExec>() | ||
| && let [child] = execution_plan.children().as_slice() |
Member
There was a problem hiding this comment.
Thanks for the fixes on the root-level case. I re-ran my repro and it is closed, and the fail-loud set_repartition_routing is exactly the right call.
One thing the simplification took with it though. Going from the multi-child walk to let [child] = execution_plan.children().as_slice() narrows this arm to single-child parents. A UnionExec no longer matches, and neither would a sort merge join range-partitioning both of its sides, which I would have guessed is one of the main consumers you have in mind for this machinery.
I ran the same probe against both commits, a UnionExec over two RuntimeStatsExec -> URRE chains:
| commit | exchanges inserted |
|---|---|
a92f0617 (before the fix) |
2 |
ea450346 (head) |
0 |
At head the plan comes out with no stage boundary anywhere:
AdaptiveDatafusionExec: is_final=false, plan_id=0, stage_id=pending, stage_resolved=false
UnionExec
RuntimeStatsExec: rows + sketch(routing=v@0 asc)
UnorderedRangeRepartitionExec: routing=v@0 asc -> 4 partitions
StatisticsExec: col_count=1, row_count=Absent
RuntimeStatsExec: rows + sketch(routing=v@0 asc)
UnorderedRangeRepartitionExec: routing=v@0 asc -> 4 partitions
StatisticsExec: col_count=1, row_count=Absent
The probe, dropped into this file's test module next to your existing ones:
#[test]
fn probe_range_repartition_under_multi_child_parent() {
let rule = DistributedExchangeRule::default();
let parent: Arc<dyn ExecutionPlan> = Arc::new(
datafusion::physical_plan::union::UnionExec::new(vec![
stats_over_urre_over_leaf(),
stats_over_urre_over_leaf(),
]),
);
let result = rule.optimize(parent, &config()).unwrap();
println!("{}", display_plan(&result));
assert_eq!(count_exchanges(result.as_ref()), 2);
}The good news is this does not go silently wrong the way the root-level case did. collect_reachable_stats only walks single-child chains, so the union stage collects no reports, and repartition_routing then errors with range-repartition stage N: no runtime-stats reports. The job fails rather than duplicating rows.
The catch is that the message points somewhere else entirely. A rule author who splices under a join input would read it as a stats collection bug rather than "your splice position is not supported yet". So either restore the multi-child walk here, or if single-child-only is a deliberate scope call for now, would you mind saying so on is_range_repartitioned so the contract is explicit? A test pinning whichever behavior you pick would be worth having too.
Investigation and repro here were LLM assisted, with the plan dumps above coming from actual runs against both commits.
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
3 participants
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Summary
I've finally done enough passes over it where I don't think it's ugly any more.
pros:
Batteries-included AQE machinery so anyone can write a
UnorderedRangeRepartitionExec/OrderedRangeRepartitionExec-inserting optimizer rule and get correct end-to-end routing for free.A range-partitioning rule now needs to do only one thing: splice a
RuntimeStatsExec+URRE/ORREabove whatever it wants range-repartitioned. Everything downstream is handled:K − 1global cuts, and rewrites the downstream location map so partitionkpulls from every producer file whose sketched[min, max]overlapsk's assigned cut range.ExchangeExec; at task-specialization time the adapter wraps theShuffleReaderin aPerPartitionFilterExec(feat(core): add PerPartitionFilterExec — per-input-partition boolean predicate #2195) that trims straddling sub-parts to their assigned slice — correctness closes without the rule ever knowing about sketches, filters, or the AQE hook.cons:
This is lots of plumbing code to achieve that. I hope I've done it in the right places. Honestly piggy backing
range_repartition_routingonExchangeExecfeels the worst. However this is following the precedent set by thecoalescefield, so maybe it's okay?Design note: the
ExchangeExecrouting slotrange_repartition_routing: Arc<Mutex<Option<RangeRepartitionRouting>>>is a direct copy of the existingcoalesce: Arc<Mutex<Option<Arc<CoalescePlan>>>>slot pattern that landed withCoalescePartitionsRule— same lifecycle (AQE rule parks decision, adapter consumes at plan-transform time), samewith_new_childrencarry-through, same idempotent-overwrite setter. That slot went through review iteration on main (b839f036 refactor(scheduler): tighten the AQE coalesce rule after review,232d7611 let the AQE coalesce slot be cleared); this PR follows the pattern rather than inventing a new one.Example consumers
Two range-repartitioning rules are being prepared as follow-ups against this infrastructure:
AdaptiveRangeShuffleRule— rewrites Hash exchanges as URRE at join-side boundaries (avoids spilling). Proof-of-concept for the machinery.ParallelWindowDetectRule(port frombrent/h2o-window-benchmarks) — wrapsBoundedWindowAggExecin URRE +RuntimeStatsExecfor parallelRANGE-frame windows. Motivating consumer for the h2o single-partition OOM case.Neither rule is in this PR — they'd conflate two decisions (the machinery vs. the specific consumer). ParallelWindowDetectRule is ready to send in short order once this lands.
Caveats
RangeRepartitionRoutingis scheduler-only state — no wire/codec changes.plan_contains_range_repartitionstops at exchange boundaries by design (a nested range repartition lives in a different stage anyway).🤖 Generated with Claude Code